"""
文档处理API路由模块
处理文档上传、解析和存储到向量数据库的功能
"""
import os
import uuid
import tempfile
from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from typing import List, Dict, Any
from app.utils.document_processor import DocumentProcessor
from app.core.qa_chain import qa_chain_system
from langchain_milvus import Milvus
from langchain_huggingface import HuggingFaceEmbeddings
from app.core.config import settings
import json

# 创建API路由器
router = APIRouter()

# 配置模板
templates = Jinja2Templates(directory="templates")

def is_valid_file_type(file: UploadFile) -> bool:
    """
    验证文件类型是否支持
    
    Args:
        file: 上传的文件
        
    Returns:
        是否为支持的文件类型
    """
    supported_types = ["text/plain", "application/pdf"]
    supported_extensions = [".txt", ".pdf"]
    
    return (file.content_type in supported_types or 
            any(file.filename.endswith(ext) for ext in supported_extensions))

def get_file_extension(file: UploadFile) -> str:
    """
    获取文件扩展名
    
    Args:
        file: 上传的文件
        
    Returns:
        文件扩展名
    """
    if file.filename.endswith(".txt"):
        return ".txt"
    elif file.filename.endswith(".pdf"):
        return ".pdf"
    else:
        return ""

def get_embeddings_model():
    """
    获取嵌入模型
    
    Returns:
        嵌入模型实例
    """
    # 初始化嵌入模型，使用BAAI/bge-large-zh-v1.5中文嵌入模型
    # 这个模型专门针对中文进行了优化，在中文文本处理方面表现更好
    model_path = "./models/bge-large-zh"
    if os.path.exists(model_path):
        # 如果本地模型存在，则使用本地模型
        # 首先修复模型文件名问题
        pooling_config_path = os.path.join(model_path, "1_Pooling", "config .json")
        fixed_pooling_config_path = os.path.join(model_path, "1_Pooling", "config.json")
        if os.path.exists(pooling_config_path) and not os.path.exists(fixed_pooling_config_path):
            os.rename(pooling_config_path, fixed_pooling_config_path)
            print("修复了模型配置文件名中的空格问题")
        
        print(f"使用本地模型: {model_path}")
        embeddings = HuggingFaceEmbeddings(
            model_name=model_path
        )
    else:
        # 否则使用在线模型（需要网络下载）
        print("使用在线模型: BAAI/bge-large-zh-v1.5")
        embeddings = HuggingFaceEmbeddings(
            model_name="BAAI/bge-large-zh-v1.5"
        )
    
    return embeddings

def store_in_vector_database(chunks: List[str], source_filename: str):
    """
    将文档块存储到向量数据库中
    
    Args:
        chunks: 文档分块列表
        source_filename: 源文件名
    """
    try:
        # 连接到Milvus向量数据库
        vector_store = Milvus(
            embedding_function=qa_chain_system.embeddings,
            collection_name="knowledge_base",
            connection_args={
                "host": qa_chain_system.milvus_host,
                "port": qa_chain_system.milvus_port
            }
        )
        
        # 为每个块添加元数据
        metadatas = [{"source": source_filename} for _ in chunks]
        
        # 批量添加到向量数据库
        vector_store.add_texts(
            texts=chunks,
            metadatas=metadatas
        )
        
        # 同时存储到Elasticsearch（如果需要）
        # 这里可以添加存储到Elasticsearch的代码
        
    except Exception as e:
        raise Exception(f"存储到向量数据库时出错: {str(e)}")

@router.get("/document_upload", response_class=HTMLResponse)
async def document_upload_page(request: Request):
    """文档上传页面"""
    return templates.TemplateResponse("document_upload.html", {"request": request})

@router.post("/upload_document")
async def upload_document(
    file: UploadFile = File(...)
):
    """
    上传并处理文档
    
    Args:
        file: 上传的文件
        
    Returns:
        处理结果
    """
    try:
        # 验证文件类型
        if not is_valid_file_type(file):
            raise HTTPException(status_code=400, detail="不支持的文件类型，请上传 .txt 或 .pdf 文件")
        
        # 创建临时文件存储上传的文件
        with tempfile.NamedTemporaryFile(delete=False, suffix=get_file_extension(file)) as temp_file:
            temp_file_path = temp_file.name
            # 将上传的文件内容写入临时文件
            content = await file.read()
            temp_file.write(content)
        
        try:
            # 处理文档
            processor = DocumentProcessor()
            
            if file.content_type == "text/plain" or file.filename.endswith(".txt"):
                # 处理文本文件
                content = processor.process_text_document(temp_file_path)
                chunks = qa_chain_system.dynamic_chunking(content)
                file_type = "text"
            elif file.content_type == "application/pdf" or file.filename.endswith(".pdf"):
                # 处理PDF文件
                result = processor.process_pdf_document(temp_file_path)
                # 提取所有页面的文本内容
                content = "\n".join([page["text"] for page in result["pages"]])
                chunks = qa_chain_system.dynamic_chunking(content)
                file_type = "pdf"
            else:
                raise HTTPException(status_code=400, detail="不支持的文件类型")
            
            # 存储到向量数据库
            store_in_vector_database(chunks, file.filename)
            
            return {
                "filename": file.filename,
                "status": "success",
                "message": f"文档已成功处理并存储到向量数据库中，共处理 {len(chunks)} 个文本块",
                "chunks_count": len(chunks)
            }
            
        finally:
            # 清理临时文件
            if os.path.exists(temp_file_path):
                os.unlink(temp_file_path)
                
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理文档时出错: {str(e)}")